{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import backtrader as bt\n",
    "import datetime\n",
    "\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# mongo feed 使用数据库作为数据源\n",
    "from mongofeed import MongoData\n",
    "\n",
    "feed = MongoData(\n",
    "    db='stock_etf',\n",
    "    dataname='etf_510300',\n",
    "    fromdate=datetime.datetime(2016,1,1),\n",
    "    todate=datetime.datetime(2019,12,31), \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class KDJ(bt.Indicator):\n",
    "    lines = ('K','D','J')\n",
    "\n",
    "    params = (\n",
    "        ('period', 9),\n",
    "        ('period_dfast', 3),\n",
    "        ('period_dslow', 3),\n",
    "    )\n",
    "    \n",
    "    plotlines = dict(\n",
    "        J=dict(\n",
    "            _fill_gt=('K', ('red', 0.50)),\n",
    "            _fill_lt=('K', ('green', 0.50)),\n",
    "        )\n",
    "    )\n",
    "    \n",
    "    def __init__(self):\n",
    "        # Add a KDJ indicator\n",
    "        self.kd = bt.indicators.StochasticFull(\n",
    "            self.data,\n",
    "            period = self.p.period,\n",
    "            period_dfast = self.p.period_dfast,\n",
    "            period_dslow = self.p.period_dslow,\n",
    "        )\n",
    "        \n",
    "        self.l.K = self.kd.percD\n",
    "        self.l.D = self.kd.percDSlow\n",
    "        self.l.J = self.K*3 - self.D*2\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class KDJStrategy(bt.Strategy):\n",
    "    params = (\n",
    "        ('period', 9),\n",
    "        ('period_dfast', 3),\n",
    "        ('period_dslow', 3),\n",
    "    )\n",
    "    \n",
    "        \n",
    "    def __init__(self):\n",
    "        \n",
    "        # use self defind a KDJ indicator\n",
    "        self.kd = KDJ(\n",
    "            self.data0,\n",
    "            period = self.p.period,\n",
    "            period_dfast = self.p.period_dfast,\n",
    "            period_dslow = self.p.period_dslow,\n",
    "        )\n",
    "    \n",
    "        self.crossover = bt.indicators.CrossOver(self.kd.K, self.kd.D, plot=False)\n",
    "        #self.above = bt.And(self.macd.macd>0.0, self.macd.macdsignal>0.0)\n",
    "        \n",
    "        self.buy_signal = (self.crossover==1)\n",
    "        self.sell_signal = (self.crossover==-1)\n",
    "        # To keep track of pending orders\n",
    "        self.order = None\n",
    "        \n",
    "        \n",
    "    def notify_order(self, order):\n",
    "        if order.status in [order.Submitted, order.Accepted]:\n",
    "            # Buy/Sell order submitted/accepted to/by broker - Nothing to do\n",
    "            return\n",
    "\n",
    "        if order.status in [order.Completed, order.Canceled, order.Margin, order.Rejected]:\n",
    "            # Write down: no pending order\n",
    "            self.order = None\n",
    "\n",
    "    def next(self):\n",
    "        # Check if an order is pending ... if yes, we cannot send a 2nd one\n",
    "        if self.order:\n",
    "            return\n",
    "        \n",
    "        # Check if we are in the market\n",
    "        if not self.position:\n",
    "            # Not yet ... we MIGHT BUY if ...\n",
    "            if self.buy_signal[0]:\n",
    "                # Keep track of the created order to avoid a 2nd order\n",
    "                self.order = self.buy()\n",
    "        else:\n",
    "            # Already in the market ... we might sell\n",
    "            if self.sell_signal[0]:\n",
    "                # Keep track of the created order to avoid a 2nd order\n",
    "                self.order = self.sell()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cerebro = bt.Cerebro()\n",
    "\n",
    "\n",
    "cerebro.adddata(feed, name= 'etf300')\n",
    "cerebro.addstrategy(KDJStrategy)\n",
    "\n",
    "\n",
    "# 小场面1万起始资金\n",
    "cerebro.broker.setcash(10000.0)\n",
    "\n",
    "# 手续费万5\n",
    "cerebro.broker.setcommission(0.0005)\n",
    "\n",
    "# 以发出信号当日收盘价成交\n",
    "cerebro.broker.set_coc(True)\n",
    "\n",
    "# Add a FixedSize sizer according to the stake\n",
    "cerebro.addsizer(bt.sizers.AllInSizerInt, percents=99)\n",
    "\n",
    "print('Starting Portfolio Value: {:.2f}'.format(cerebro.broker.getvalue()))\n",
    "\n",
    "cerebro.addanalyzer(bt.analyzers.SQN)\n",
    "\n",
    "result = cerebro.run()\n",
    "\n",
    "print('Ending Portfolio Value: {:.2f}'.format(cerebro.broker.getvalue()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ana = result[0].analyzers.sqn.get_analysis()\n",
    "print(\"sqn: {:.3f}, trades:{:d}\".format(ana['sqn'],ana['trades']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cerebro.plot(iplot=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
